// 支持大数组、多 block 的 CUDA 基数排序（基于 CUB, thrust）
#include <iostream>
#include <algorithm>
#include <cuda_runtime.h>
#include <chrono>
#include <vector>
#include <cub/cub.cuh>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/sequence.h>

using namespace std;
using namespace std::chrono;

const int num_data = 1 << 20;  // 1048576 ≈ 1M 个 unsigned int


// 基数排序CPU优化版本
void cpu_radix_sort(unsigned int* arr, int n) {
    unsigned int *buffer = new unsigned int[n];
    unsigned int *input = arr;
    unsigned int *output = buffer;

    for (int bit = 0; bit < 32; bit++) {
        int count[2] = {0, 0};

        // 统计
        for (int i = 0; i < n; i++) {
            int digit = (input[i] >> bit) & 1;
            count[digit]++;
        }

        // 起始位置
        int offset[2] = {0, count[0]};

        // Scatter
        for (int i = 0; i < n; i++) {
            int digit = (input[i] >> bit) & 1;
            output[offset[digit]++] = input[i];
        }

        std::swap(input, output);
    }

    // 结果在 input 中
    if (input != arr) {
        std::copy(input, input + n, arr);
    }

    delete[] buffer;
}

void cpu_radix_sort_with_index(unsigned int* arr, unsigned int* arr_pos,
                               unsigned int* sorted_arr, unsigned int* original_indices,
                               int n) {
    // 分配缓冲区：用于 values 和 positions 的双缓冲
    unsigned int* values_curr = new unsigned int[n];
    unsigned int* values_next = new unsigned int[n];
    unsigned int* pos_curr = new unsigned int[n];
    unsigned int* pos_next = new unsigned int[n];

    // 初始化当前值和位置数组
    std::copy(arr, arr + n, values_curr);
    std::copy(arr_pos, arr_pos + n, pos_curr);

    for (int bit = 0; bit < 32; bit++) {
        int count[2] = {0, 0};

        // 统计当前 bit 位上 0 和 1 的数量
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            count[digit]++;
        }

        // 偏移：0 放前面，1 紧随其后
        int offset[2] = {0, count[0]};

        // Scatter 阶段：根据当前 bit 位，将值和位置同步写入 next 数组
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            int dest_idx = offset[digit]++;
            values_next[dest_idx] = values_curr[i];
            pos_next[dest_idx] = pos_curr[i];
        }

        // 交换 curr 和 next（通过指针交换，避免数据拷贝）
        std::swap(values_curr, values_next);
        std::swap(pos_curr, pos_next);
    }

    // 此时结果在 values_curr 和 pos_curr 中
    std::copy(values_curr, values_curr + n, sorted_arr);
    std::copy(pos_curr, pos_curr + n, original_indices);

    // 释放动态内存
    delete[] values_curr;
    delete[] values_next;
    delete[] pos_curr;
    delete[] pos_next;
}

struct CompareTuple {
    __device__ 
    bool operator()(const thrust::tuple<unsigned int, unsigned int>& a,
                   const thrust::tuple<unsigned int, unsigned int>& b) {
        return thrust::get<0>(a) < thrust::get<0>(b);
    }
};

void thrust_radix_sort(unsigned int* values, unsigned int* indices, 
                      unsigned int* out, unsigned int* out_pos, int n) {
    // 创建设备向量
    thrust::device_vector<unsigned int> d_values(values, values + n);
    thrust::device_vector<unsigned int> d_indices(n);
    
    // 生成初始索引序列
    thrust::sequence(d_indices.begin(), d_indices.end());
    
    // 创建zip迭代器，将值和索引绑定
    auto begin = thrust::make_zip_iterator(
        thrust::make_tuple(d_values.begin(), d_indices.begin()));
    auto end = thrust::make_zip_iterator(
        thrust::make_tuple(d_values.end(), d_indices.end()));
    
    // 按值排序，同时保持索引关联
    thrust::sort(begin, end, CompareTuple());
    
    // 将结果拷贝回主机
    thrust::copy(d_values.begin(), d_values.end(), out);
    thrust::copy(d_indices.begin(), d_indices.end(), out_pos);
}

void cub_radix_sort(unsigned int* values, unsigned int* indices, 
                    unsigned int* out, unsigned int* out_pos, int n){
    void* d_temp_storage = nullptr;
    size_t temp_storage_bytes = 0;

    // 查询临时存储大小
    cub::DeviceRadixSort::SortPairs(
        d_temp_storage, temp_storage_bytes,
        values, out,
        indices, out_pos,
        num_data
    );

    // 分配临时内存
    cudaMalloc(&d_temp_storage, temp_storage_bytes);

    cub::DeviceRadixSort::SortPairs(
        d_temp_storage, temp_storage_bytes,
        values, out,
        indices, out_pos,
        num_data
    );
    cudaFree(d_temp_storage);
}

int main() {
    // Host 内存
    unsigned int* h_data = new unsigned int[num_data];
    unsigned int* cpu_data = new unsigned int[num_data];
    unsigned int* cpu_data_index = new unsigned int[num_data];
    unsigned int* cpu_out = new unsigned int[num_data];
    unsigned int* cpu_out_index = new unsigned int[num_data];
    unsigned int* h_index = new unsigned int[num_data];  // 存储排序后每个元素的原始索引

    // 初始化随机数据
    srand(time(0));
    for (int i = 0; i < num_data; ++i) {
        h_data[i] = rand() | (static_cast<unsigned int>(rand()) << 15);
    }
    memcpy(cpu_data, h_data, sizeof(unsigned int) * num_data);
    // 初始化索引数组
    for (int i = 0; i < num_data; ++i) {
        h_index[i] = i;
        cpu_data_index[i] = i;
    }
    
    // GPU 内存
    unsigned int *d_keys_in, *d_keys_out;
    unsigned int *d_values_in, *d_values_out;  // values 用于存储原始索引

    cudaMalloc(&d_keys_in, sizeof(unsigned int) * num_data);
    cudaMalloc(&d_keys_out, sizeof(unsigned int) * num_data);
    cudaMalloc(&d_values_in, sizeof(unsigned int) * num_data);
    cudaMalloc(&d_values_out, sizeof(unsigned int) * num_data);
    cudaMemcpy(d_keys_in, h_data, sizeof(unsigned int) * num_data, cudaMemcpyHostToDevice);
    cudaMemcpy(d_values_in, h_index, sizeof(unsigned int) * num_data, cudaMemcpyHostToDevice);

    // 开始计时
    auto start = high_resolution_clock::now();
    // thrust_radix_sort(d_keys_in, d_values_in, d_keys_out, d_values_out, num_data);
    cub_radix_sort(d_keys_in, d_values_in, d_keys_out, d_values_out, num_data);
    cudaDeviceSynchronize();
    auto end = high_resolution_clock::now();
    auto gpu_time = duration_cast<microseconds>(end - start).count();

    // 拷贝结果
    cudaMemcpy(h_data, d_keys_out, sizeof(unsigned int) * num_data, cudaMemcpyDeviceToHost);
    cudaMemcpy(h_index, d_values_out, sizeof(unsigned int) * num_data, cudaMemcpyDeviceToHost);

    // ========================
    // CPU 排序（带索引）
    // ========================
    vector<pair<unsigned int, unsigned int>> cpu_pairs(num_data);
    for (int i = 0; i < num_data; ++i) {
        cpu_pairs[i] = make_pair(cpu_data[i], i);
    }

    auto start_cpu = high_resolution_clock::now();
    cpu_radix_sort_with_index(cpu_data, cpu_data_index, cpu_out, cpu_out_index, num_data);

    auto end_cpu = high_resolution_clock::now();
    auto cpu_time = duration_cast<microseconds>(end_cpu - start_cpu).count();

    // 提取 CPU 排序后的值和索引用于验证
    vector<unsigned int> cpu_sorted_values(num_data);
    vector<unsigned int> cpu_sorted_indices(num_data);
    for (int i = 0; i < num_data; ++i) {
        cpu_sorted_values[i] = cpu_pairs[i].first;
        cpu_sorted_indices[i] = cpu_pairs[i].second;
    }

    // ========================
    // 验证结果（值 + 索引）
    // ========================
    bool is_correct = true;
    for (int i = 0; i < num_data; ++i) {
        if (h_data[i] != cpu_out[i]) {
            std::cout << "Value mismatch at index " << i
                      << ": GPU=" << h_data[i] << ", CPU=" << cpu_out[i] << std::endl;
            is_correct = false;
            break;
        }
        if (h_index[i] != cpu_out_index[i]) {
            std::cout << "Index mismatch at position " << i
                      << ": GPU_index=" << h_index[i] << ", CPU_index=" << cpu_out_index[i] << std::endl;
            is_correct = false;
            break;
        }
    }
    if (is_correct) {
        std::cout << "Result is CORRECT!" << std::endl;
    } else {
        std::cout << "Result is INCORRECT!" << std::endl;
    }

    // ========================
    // 输出性能
    // ========================
    std::cout << "Data size: " << num_data << " elements (" << (num_data * 4.0 / (1024*1024)) << " MB)" << std::endl;
    std::cout << "GPU  time: " << gpu_time << " μs (" << gpu_time / 1000.0 << " ms)" << std::endl;
    std::cout << "CPU  time: " << cpu_time << " μs (" << cpu_time / 1000.0 << " ms)" << std::endl;
    if (cpu_time > 0) {
        std::cout << "Speedup: " << (double)cpu_time / gpu_time << "x" << std::endl;
    }

    // ========================
    // 释放内存
    // ========================
    delete[] h_data;
    delete[] cpu_data;
    delete[] cpu_data_index;
    delete[] cpu_out;
    delete[] cpu_out_index;
    delete[] h_index;
    cudaFree(d_keys_in);
    cudaFree(d_keys_out);
    cudaFree(d_values_in);
    cudaFree(d_values_out);
    
    return 0;
}